{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "origin_pos": 0
   },
   "source": [
    "# Multi-Head Attention\n",
    ":label:`sec_multihead-attention`\n",
    "\n",
    "\n",
    "In practice,\n",
    "given the same set of queries, keys, and values\n",
    "we may want our model to\n",
    "combine knowledge from\n",
    "different behaviors of the same attention mechanism,\n",
    "such as capturing dependencies of various ranges (e.g., shorter-range vs. longer-range)\n",
    "within a sequence.\n",
    "Thus, \n",
    "it may be beneficial \n",
    "to allow our attention mechanism\n",
    "to jointly use different representation subspaces\n",
    "of queries, keys, and values.\n",
    "\n",
    "\n",
    "\n",
    "To this end,\n",
    "instead of performing a single attention pooling,\n",
    "queries, keys, and values\n",
    "can be transformed\n",
    "with $h$ independently learned linear projections.\n",
    "Then these $h$ projected queries, keys, and values\n",
    "are fed into attention pooling in parallel.\n",
    "In the end,\n",
    "$h$ attention pooling outputs\n",
    "are concatenated and \n",
    "transformed with another learned linear projection\n",
    "to produce the final output.\n",
    "This design\n",
    "is called *multi-head attention*,\n",
    "where each of the $h$ attention pooling outputs\n",
    "is a *head* :cite:`Vaswani.Shazeer.Parmar.ea.2017`.\n",
    "Using fully-connected layers\n",
    "to perform learnable linear transformations,\n",
    ":numref:`fig_multi-head-attention`\n",
    "describes multi-head attention.\n",
    "\n",
    "![Multi-head attention, where multiple heads are concatenated then linearly transformed.](https://raw.githubusercontent.com/d2l-ai/d2l-en/master/img/multi-head-attention.svg)\n",
    ":label:`fig_multi-head-attention`\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "## Model\n",
    "\n",
    "Before providing the implementation of multi-head attention,\n",
    "let us formalize this model mathematically.\n",
    "Given a query $\\mathbf{q} \\in \\mathbb{R}^{d_q}$,\n",
    "a key $\\mathbf{k} \\in \\mathbb{R}^{d_k}$,\n",
    "and a value $\\mathbf{v} \\in \\mathbb{R}^{d_v}$,\n",
    "each attention head $\\mathbf{h}_i$  ($i = 1, \\ldots, h$)\n",
    "is computed as\n",
    "\n",
    "$$\\mathbf{h}_i = f(\\mathbf W_i^{(q)}\\mathbf q, \\mathbf W_i^{(k)}\\mathbf k,\\mathbf W_i^{(v)}\\mathbf v) \\in \\mathbb R^{p_v},$$\n",
    "\n",
    "where learnable parameters\n",
    "$\\mathbf W_i^{(q)}\\in\\mathbb R^{p_q\\times d_q}$,\n",
    "$\\mathbf W_i^{(k)}\\in\\mathbb R^{p_k\\times d_k}$\n",
    "and $\\mathbf W_i^{(v)}\\in\\mathbb R^{p_v\\times d_v}$,\n",
    "and\n",
    "$f$ is attention pooling,\n",
    "such as\n",
    "additive attention and scaled dot-product attention\n",
    "in :numref:`sec_attention-scoring-functions`.\n",
    "The multi-head attention output\n",
    "is another linear transformation via \n",
    "learnable parameters\n",
    "$\\mathbf W_o\\in\\mathbb R^{p_o\\times h p_v}$\n",
    "of the concatenation of $h$ heads:\n",
    "\n",
    "$$\\mathbf W_o \\begin{bmatrix}\\mathbf h_1\\\\\\vdots\\\\\\mathbf h_h\\end{bmatrix} \\in \\mathbb{R}^{p_o}.$$\n",
    "\n",
    "Based on this design,\n",
    "each head may attend to different parts of the input.\n",
    "More sophisticated functions than the simple weighted average\n",
    "can be expressed.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load ../utils/djl-imports\n",
    "%load ../utils/plot-utils\n",
    "%load ../utils/Functions.java\n",
    "%load ../utils/PlotUtils.java\n",
    "\n",
    "%load ../utils/attention/Chap10Utils.java\n",
    "%load ../utils/attention/DotProductAttention.java\n",
    "%load ../utils/attention/MultiHeadAttention.java\n",
    "%load ../utils/attention/PositionalEncoding.java"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "NDManager manager = NDManager.newBaseManager();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "origin_pos": 6
   },
   "source": [
    "To allow for parallel computation of multiple heads,\n",
    "the below `MultiHeadAttention` class uses two transposition functions as defined below.\n",
    "Specifically,\n",
    "the `transposeOutput` function reverses the operation\n",
    "of the `transposeQkv` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "public static NDArray transposeQkv(NDArray X, int numHeads) {\n",
    "    // Shape of input `X`:\n",
    "    // (`batchSize`, no. of queries or key-value pairs, `numHiddens`).\n",
    "    // Shape of output `X`:\n",
    "    // (`batchSize`, no. of queries or key-value pairs, `numHeads`,\n",
    "    // `numHiddens` / `numHeads`)\n",
    "    X = X.reshape(X.getShape().get(0), X.getShape().get(1), numHeads, -1);\n",
    "\n",
    "    // Shape of output `X`:\n",
    "    // (`batchSize`, `numHeads`, no. of queries or key-value pairs,\n",
    "    // `numHiddens` / `numHeads`)\n",
    "    X = X.transpose(0, 2, 1, 3);\n",
    "\n",
    "    // Shape of `output`:\n",
    "    // (`batchSize` * `numHeads`, no. of queries or key-value pairs,\n",
    "    // `numHiddens` / `numHeads`)\n",
    "    return X.reshape(-1, X.getShape().get(2), X.getShape().get(3));\n",
    "}\n",
    "\n",
    "public static NDArray transposeOutput(NDArray X, int numHeads) {\n",
    "    X = X.reshape(-1, numHeads, X.getShape().get(1), X.getShape().get(2));\n",
    "    X = X.transpose(0, 2, 1, 3);\n",
    "    return X.reshape(X.getShape().get(0), X.getShape().get(1), -1);\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "origin_pos": 3
   },
   "source": [
    "## Implementation\n",
    "\n",
    "In our implementation,\n",
    "we choose the scaled dot-product attention\n",
    "for each head of the multi-head attention.\n",
    "To avoid significant growth\n",
    "of computational cost and parameterization cost,\n",
    "we set\n",
    "$p_q = p_k = p_v = p_o / h$.\n",
    "Note that $h$ heads\n",
    "can be computed in parallel\n",
    "if we set\n",
    "the number of outputs of linear transformations\n",
    "for the query, key, and value\n",
    "to $p_q h = p_k h = p_v h = p_o$.\n",
    "In the following implementation,\n",
    "$p_o$ is specified via the argument `numHiddens`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "public static class MultiHeadAttention extends AbstractBlock {\n",
    "\n",
    "    private int numHeads;\n",
    "    public DotProductAttention attention;\n",
    "    private Linear W_k;\n",
    "    private Linear W_q;\n",
    "    private Linear W_v;\n",
    "    private Linear W_o;\n",
    "\n",
    "    public MultiHeadAttention(int numHiddens, int numHeads, float dropout, boolean useBias) {\n",
    "        this.numHeads = numHeads;\n",
    "\n",
    "        attention = new DotProductAttention(dropout);\n",
    "\n",
    "        W_q = Linear.builder().setUnits(numHiddens).optBias(useBias).build();\n",
    "        addChildBlock(\"W_q\", W_q);\n",
    "\n",
    "        W_k = Linear.builder().setUnits(numHiddens).optBias(useBias).build();\n",
    "        addChildBlock(\"W_k\", W_k);\n",
    "\n",
    "        W_v = Linear.builder().setUnits(numHiddens).optBias(useBias).build();\n",
    "        addChildBlock(\"W_v\", W_v);\n",
    "\n",
    "        W_o = Linear.builder().setUnits(numHiddens).optBias(useBias).build();\n",
    "        addChildBlock(\"W_o\", W_o);\n",
    "\n",
    "        Dropout dropout1 = Dropout.builder().optRate(dropout).build();\n",
    "        addChildBlock(\"dropout\", dropout1);\n",
    "    }\n",
    "\n",
    "    @Override\n",
    "    protected NDList forwardInternal(\n",
    "            ParameterStore ps,\n",
    "            NDList inputs,\n",
    "            boolean training,\n",
    "            PairList<String, Object> params) {\n",
    "        // Shape of `queries`, `keys`, or `values`:\n",
    "        // (`batchSize`, no. of queries or key-value pairs, `numHiddens`)\n",
    "        // Shape of `validLens`:\n",
    "        // (`batchSize`,) or (`batchSize`, no. of queries)\n",
    "        // After transposing, shape of output `queries`, `keys`, or `values`:\n",
    "        // (`batchSize` * `numHeads`, no. of queries or key-value pairs,\n",
    "        // `numHiddens` / `numHeads`)\n",
    "        NDArray queries = inputs.get(0);\n",
    "        NDArray keys = inputs.get(1);\n",
    "        NDArray values = inputs.get(2);\n",
    "        NDArray validLens = inputs.get(3);\n",
    "        // On axis 0, copy the first item (scalar or vector) for\n",
    "        // `numHeads` times, then copy the next item, and so on\n",
    "        validLens = validLens.repeat(0, numHeads);\n",
    "\n",
    "        queries =\n",
    "                transposeQkv(\n",
    "                        W_q.forward(ps, new NDList(queries), training, params).get(0),\n",
    "                        numHeads);\n",
    "        keys =\n",
    "                transposeQkv(\n",
    "                        W_k.forward(ps, new NDList(keys), training, params).get(0), numHeads);\n",
    "        values =\n",
    "                transposeQkv(\n",
    "                        W_v.forward(ps, new NDList(values), training, params).get(0), numHeads);\n",
    "\n",
    "        // Shape of `output`: (`batchSize` * `numHeads`, no. of queries,\n",
    "        // `numHiddens` / `numHeads`)\n",
    "        NDArray output =\n",
    "                attention\n",
    "                        .forward(\n",
    "                                ps,\n",
    "                                new NDList(queries, keys, values, validLens),\n",
    "                                training,\n",
    "                                params)\n",
    "                        .get(0);\n",
    "\n",
    "        // Shape of `outputConcat`:\n",
    "        // (`batchSize`, no. of queries, `numHiddens`)\n",
    "        NDArray outputConcat = transposeOutput(output, numHeads);\n",
    "        return new NDList(W_o.forward(ps, new NDList(outputConcat), training, params).get(0));\n",
    "    }\n",
    "\n",
    "    @Override\n",
    "    public Shape[] getOutputShapes(Shape[] inputShapes) {\n",
    "        throw new UnsupportedOperationException(\"Not implemented\");\n",
    "    }\n",
    "\n",
    "    @Override\n",
    "    public void initializeChildBlocks(\n",
    "            NDManager manager, DataType dataType, Shape... inputShapes) {\n",
    "        try (NDManager sub = manager.newSubManager()) {\n",
    "            NDArray queries = sub.zeros(inputShapes[0], dataType);\n",
    "            NDArray keys = sub.zeros(inputShapes[1], dataType);\n",
    "            NDArray values = sub.zeros(inputShapes[2], dataType);\n",
    "            NDArray validLens = sub.zeros(inputShapes[3], dataType);\n",
    "            validLens = validLens.repeat(0, numHeads);\n",
    "\n",
    "            ParameterStore ps = new ParameterStore(sub, false);\n",
    "\n",
    "            W_q.initialize(manager, dataType, queries.getShape());\n",
    "            W_k.initialize(manager, dataType, keys.getShape());\n",
    "            W_v.initialize(manager, dataType, values.getShape());\n",
    "\n",
    "            queries =\n",
    "                    transposeQkv(W_q.forward(ps, new NDList(queries), false).get(0), numHeads);\n",
    "            keys = transposeQkv(W_k.forward(ps, new NDList(keys), false).get(0), numHeads);\n",
    "            values = transposeQkv(W_v.forward(ps, new NDList(values), false).get(0), numHeads);\n",
    "\n",
    "            NDList list = new NDList(queries, keys, values, validLens);\n",
    "            attention.initialize(sub, dataType, list.getShapes());\n",
    "            NDArray output = attention.forward(ps, list, false).head();\n",
    "            NDArray outputConcat = Chap10Utils.transposeOutput(output, numHeads);\n",
    "\n",
    "            W_o.initialize(manager, dataType, outputConcat.getShape());\n",
    "        }\n",
    "    }\n",
    "}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "origin_pos": 9
   },
   "source": [
    "Let us test our implemented `MultiHeadAttention` class\n",
    "using a toy example where keys and values are the same.\n",
    "As a result,\n",
    "the shape of the multi-head attention output\n",
    "is (`batchSize`, `numQueries`, `numHiddens`).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "int numHiddens = 100;\n",
    "int numHeads = 5;\n",
    "MultiHeadAttention attention = new MultiHeadAttention(numHiddens, numHeads, 0.5f, false);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "int batchSize = 2;\n",
    "int numQueries = 4;\n",
    "int numKvpairs = 6;\n",
    "NDArray validLens = manager.create(new float[] {3, 2});\n",
    "NDArray X = manager.ones(new Shape(batchSize, numQueries, numHiddens));\n",
    "NDArray Y = manager.ones(new Shape(batchSize, numKvpairs, numHiddens));\n",
    "\n",
    "ParameterStore ps = new ParameterStore(manager, false);\n",
    "NDList input = new NDList(X, Y, Y, validLens);\n",
    "attention.initialize(manager, DataType.FLOAT32, input.getShapes());\n",
    "NDList result = attention.forward(ps, input, false);\n",
    "result.get(0).getShape();\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "origin_pos": 13
   },
   "source": [
    "## Summary\n",
    "\n",
    "* Multi-head attention combines knowledge of the same attention pooling via different representation subspaces of queries, keys, and values.\n",
    "* To compute multiple heads of multi-head attention in parallel, proper tensor manipulation is needed.\n",
    "\n",
    "\n",
    "\n",
    "## Exercises\n",
    "\n",
    "1. Visualize attention weights of multiple heads in this experiment.\n",
    "1. Suppose that we have a trained model based on multi-head attention and we want to prune least important attention heads to increase the prediction speed. How can we design experiments to measure the importance of an attention head?\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Java",
   "language": "java",
   "name": "java"
  },
  "language_info": {
   "codemirror_mode": "java",
   "file_extension": ".jshell",
   "mimetype": "text/x-java-source",
   "name": "Java",
   "pygments_lexer": "java",
   "version": "14.0.2+12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
